﻿#pragma once
#include<iostream>
#include<opencv2/opencv.hpp>
#define YOLO_P6 false // Whether to use P6 Model //
struct Output {
    int id; // Result category id/
    float confidence; // Result confidence //
    cv::Rect box; // Rectangle box //
};
class Yolo {
public:
    Yolo() {
    }
    ~Yolo() {}
    bool readModel(cv::dnn::Net& net, std::string& netPath, bool isCuda);
    bool Detect(cv::Mat& SrcImg, cv::dnn::Net& net, std::vector<Output>& output,std::vector<std::string>& className);
    void drawPred(cv::Mat& img, std::vector<Output> result, std::vector<cv::Scalar> color,std::vector<std::string>& className);
    const int netWidth = 960; //ONNX Picture input width yolov5s.onnx 960
    const int netHeight = 960; //ONNX Picture input height
private:
#if(defined YOLO_P6 && YOLO_P6==true)
    const float netAnchors[4][6] = { { 19,27, 44,40, 38,94 },{ 96,68, 86,152, 180,137 },{ 140,301, 303,264, 238,542 },{ 436,615, 739,380, 925,792 } };
    const int netWidth = 1280; //ONNX Picture input width
    const int netHeight = 1280; //ONNX Picture input height
    const int strideSize = 4; //stride size
#else
    const float netAnchors[3][6] = { { 10,13, 16,30, 33,23 },{ 30,61, 62,45, 59,119 },{ 116,90, 156,198, 373,326 } };
    const int strideSize = 3; //stride size
#endif // YOLO_P6
    const float netStride[4] = { 8, 16.0,32,64 };
    float boxThreshold = 0.25;
    float classThreshold = 0.25;
    float nmsThreshold = 0.45;
    float nmsScoreThreshold = boxThreshold * classThreshold;
};
